#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
连接池模块 - 管理 TCP 连接池
"""

import socket
import queue
import threading
import errno
import time
from logger import setup_logger
from config import get_config_value

# 设置日志记录器
logger = setup_logger("connection_pool")

class ConnectionPool:
    """TCP 连接池，用于管理和重用远程连接"""
    def __init__(self, remote_host, remote_port):
        """初始化连接池
        
        Args:
            remote_host: 远程主机地址
            remote_port: 远程主机端口
        """
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.max_connections = get_config_value("tcp", "max_connections", 50)
        self.min_connections = get_config_value("tcp", "min_connections", 10)
        self.timeout = get_config_value("tcp", "timeout", 300)
        self.pool = queue.Queue(maxsize=self.max_connections)
        self.active_connections = 0  # 当前活跃连接数
        self.lock = threading.Lock()
        self.last_resize_time = time.time()
        self.resize_interval = 30  # 调整大小的时间间隔（秒）
        self.usage_history = []  # 存储使用率历史
        self.connection_last_used = {}  # 存储连接最后使用时间
        self.idle_timeout = 600  # 空闲连接超时时间（秒）
        
        # 初始化连接池
        self._initialize_pool(self.min_connections)
        
        # 启动监控线程
        self.monitor_thread = threading.Thread(target=self._monitor_pool, daemon=True)
        self.monitor_thread.start()
    
    def _initialize_pool(self, num_connections):
        """初始化连接池"""
        for _ in range(num_connections):
            try:
                conn = self._create_connection()
                self.pool.put(conn)
            except Exception as e:
                logger.error(f"初始化连接池时创建连接失败: {e}")
    
    def _monitor_pool(self):
        """监控连接池并动态调整大小"""
        while True:
            try:
                current_time = time.time()
                # 每隔一段时间检查一次
                if current_time - self.last_resize_time > self.resize_interval:
                    self._resize_pool()
                    self.last_resize_time = current_time
                
                # 检查空闲连接是否需要关闭
                self._cleanup_idle_connections()
                
                # 记录当前使用率
                with self.lock:
                    usage = 1.0 - (self.pool.qsize() / max(self.max_connections, 1))
                    self.usage_history.append(usage)
                    # 只保留最近10次的使用率记录
                    if len(self.usage_history) > 10:
                        self.usage_history.pop(0)
                
                time.sleep(5)  # 每5秒检查一次
            except Exception as e:
                logger.error(f"连接池监控线程出错: {e}")
                time.sleep(10)  # 出错后等待时间更长
    
    def _resize_pool(self):
        """根据使用率动态调整连接池大小"""
        if not self.usage_history:
            return
        
        # 计算平均使用率
        avg_usage = sum(self.usage_history) / len(self.usage_history)
        
        with self.lock:
            current_size = self.pool.qsize() + self.active_connections
            
            # 高负载情况 (使用率 > 80%)，增加连接数
            if avg_usage > 0.8 and current_size < self.max_connections:
                new_connections = min(5, self.max_connections - current_size)
                logger.info(f"连接池负载较高 ({avg_usage:.2f})，增加 {new_connections} 个连接")
                for _ in range(new_connections):
                    try:
                        conn = self._create_connection()
                        self.pool.put(conn)
                    except Exception as e:
                        logger.error(f"增加连接失败: {e}")
            
            # 低负载情况 (使用率 < 30%)，减少连接数
            elif avg_usage < 0.3 and self.pool.qsize() > self.min_connections:
                reduce_count = min(3, self.pool.qsize() - self.min_connections)
                logger.info(f"连接池负载较低 ({avg_usage:.2f})，减少 {reduce_count} 个连接")
                for _ in range(reduce_count):
                    try:
                        conn = self.pool.get(block=False)
                        conn.close()
                    except queue.Empty:
                        break
                    except Exception as e:
                        logger.error(f"减少连接失败: {e}")
    
    def _cleanup_idle_connections(self):
        """清理空闲连接"""
        try:
            current_time = time.time()
            idle_connections = []
            
            # 检查连接池中的空闲连接
            with self.lock:
                # 如果连接池大小小于等于最小连接数，不进行清理
                if self.pool.qsize() <= self.min_connections:
                    return
                
                # 获取所有连接并检查空闲时间
                temp_queue = queue.Queue()
                idle_count = 0
                
                while not self.pool.empty():
                    try:
                        conn = self.pool.get(block=False)
                        conn_id = id(conn)
                        
                        # 检查连接是否空闲超时
                        last_used_time = self.connection_last_used.get(conn_id, current_time)
                        idle_time = current_time - last_used_time
                        
                        if idle_time > self.idle_timeout and self.pool.qsize() + idle_count + 1 > self.min_connections:
                            # 连接空闲时间超过阈值且关闭后连接数仍大于最小值，关闭连接
                            try:
                                conn.close()
                                if conn_id in self.connection_last_used:
                                    del self.connection_last_used[conn_id]
                                idle_count += 1
                                logger.debug(f"关闭空闲连接，空闲时间: {idle_time:.1f}秒")
                            except Exception as e:
                                logger.error(f"关闭空闲连接时出错: {e}")
                                # 关闭失败，放回池中
                                temp_queue.put(conn)
                        else:
                            # 连接未超时或连接数已达最小值，放回池中
                            temp_queue.put(conn)
                    except queue.Empty:
                        break
                
                # 将剩余连接放回原池
                while not temp_queue.empty():
                    try:
                        self.pool.put(temp_queue.get(block=False))
                    except queue.Empty:
                        break
                
                if idle_count > 0:
                    logger.info(f"已清理 {idle_count} 个空闲连接，当前连接池大小: {self.pool.qsize()}")
        except Exception as e:
            logger.error(f"清理空闲连接时出错: {e}")
    
    def _create_connection(self):
        """创建新的连接，带重试机制"""
        max_retries = 3
        retry_delay = 1
        
        for attempt in range(max_retries):
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 262144)  # 增加到256KB
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 262144)  # 增加到256KB
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                
                #移除不支持的TCP窗口大小设置
                #sock.setsockopt(socket.SOL_TCP, socket.TCP_WINDOW_CLAMP, 262144)  # 设置最大窗口大小为256KB
                
                # 在 Windows 上设置 TCP keepalive 参数
                sock.ioctl(socket.SIO_KEEPALIVE_VALS, (1, 60000, 5000))
                
                # 设置连接超时
                sock.settimeout(10)
                
                # 连接到远程服务器
                sock.connect((self.remote_host, self.remote_port))
                
                # 连接成功后恢复非阻塞模式
                sock.setblocking(False)
                logger.debug(f"连接池: 成功创建到 {self.remote_host}:{self.remote_port} 的新连接")
                return sock
            except Exception as e:
                if attempt < max_retries - 1:
                    logger.warning(f"连接失败，尝试重新连接 ({attempt+1}/{max_retries}): {e}")
                    time.sleep(retry_delay)
                    retry_delay *= 2  # 指数退避
                else:
                    logger.error(f"连接到 {self.remote_host}:{self.remote_port} 失败，已达到最大重试次数: {e}")
                    raise
    
    def cleanup(self):
        """清理所有连接"""
        logger.info("连接池: 开始清理所有连接")
        closed_count = 0
        
        while not self.pool.empty():
            try:
                conn = self.pool.get(block=False)
                conn.close()
                closed_count += 1
                # 不应该减少活跃连接计数，因为池中的连接不是活跃的
            except queue.Empty:
                break
            except Exception as e:
                logger.error(f"关闭连接时出错: {e}")
        
        logger.info(f"连接池: 已清理 {closed_count} 个连接")


    def get_connection(self):
        """获取连接，增加重试和健康检查机制"""
        start_time = time.time()
        retry_count = 0
        max_retries = 3
        
        while retry_count < max_retries:
            try:
                # 尝试从池中获取连接
                conn = self.pool.get(block=True, timeout=self.timeout)
                
                # 检查连接是否健康
                if self._is_connection_healthy(conn):
                    with self.lock:
                        self.active_connections += 1
                        # 更新连接最后使用时间
                        self.connection_last_used[id(conn)] = time.time()
                    return conn
                else:
                    # 连接不健康，关闭并创建新连接
                    logger.warning("获取到不健康的连接，正在替换")
                    try:
                        conn.close()
                    except:
                        pass
                    
                    # 创建新连接
                    new_conn = self._create_connection()
                    with self.lock:
                        self.active_connections += 1
                        # 更新连接最后使用时间
                        self.connection_last_used[id(new_conn)] = time.time()  # 使用 new_conn
                    return new_conn
            
            except queue.Empty:
                # 连接池为空，尝试创建新连接
                if self.active_connections + self.pool.qsize() < self.max_connections:
                    try:
                        logger.info("连接池已空，创建新连接")
                        # 创建新连接
                        new_conn = self._create_connection()
                        with self.lock:
                            self.active_connections += 1
                        return new_conn
                    except Exception as e:
                        logger.error(f"创建新连接失败: {e}",exc_info=True)
                
                # 如果已达到最大连接数，等待并重试
                retry_count += 1
                logger.warning(f"连接池: 等待连接超时 (重试 {retry_count}/{max_retries})")
                
                # 如果等待时间过长，尝试关闭一些空闲时间最长的连接
                if time.time() - start_time > self.timeout * 2:
                    self._force_recycle_connections()
        
        # 所有重试都失败
        raise Exception("无法获取连接: 连接池已满且等待超时")
    
    def _is_connection_healthy(self, conn):
        """检查连接是否健康"""
        try:
            # 设置非阻塞模式进行检查
            orig_blocking = conn.getblocking()
            conn.setblocking(False)
            
            # 尝试发送0字节数据作为健康检查
            conn.send(b'')
            
            # 尝试接收数据，但不期望有数据返回
            try:
                conn.recv(1)
            except (socket.error, OSError) as e:
                if isinstance(e, socket.error) and e.errno in (errno.EWOULDBLOCK, errno.EAGAIN):
                    # 这些错误表示连接仍然有效，只是暂时没有数据
                    pass
                else:
                    # 其他错误表示连接可能已断开
                    conn.setblocking(orig_blocking)
                    return False
            
            # 恢复原始阻塞模式
            conn.setblocking(orig_blocking)
            return True
        except (socket.error, OSError) as e:
            if isinstance(e, socket.error) and e.errno in (errno.EWOULDBLOCK, errno.EAGAIN):
                # 这些错误表示连接仍然有效，只是暂时不可写
                conn.setblocking(orig_blocking)
                return True
            logger.warning(f"连接健康检查失败: {e}")
            return False
        except Exception as e:
            logger.warning(f"连接健康检查出现未知错误: {e}")
            return False
    
    def _force_recycle_connections(self):
        """强制回收一些连接以缓解连接池压力"""
        logger.warning("强制回收连接以缓解连接池压力")
        try:
            with self.lock:
                # 如果连接池为空，无法回收
                if self.pool.empty():
                    logger.warning("连接池为空，无法回收连接")
                    return
                
                # 回收连接数量，回收约20%的连接，至少1个
                recycle_count = max(1, int(self.pool.qsize() * 0.2))
                logger.info(f"尝试回收 {recycle_count} 个连接")
                
                # 获取所有连接
                connections = []
                temp_queue = queue.Queue()
                
                while not self.pool.empty():
                    try:
                        conn = self.pool.get(block=False)
                        conn_id = id(conn)
                        last_used_time = self.connection_last_used.get(conn_id, 0)
                        connections.append((conn, last_used_time, conn_id))
                    except queue.Empty:
                        break
                
                # 按最后使用时间排序，优先关闭最长时间未使用的连接
                connections.sort(key=lambda x: x[1])
                
                # 关闭部分连接
                closed_count = 0
                for i, (conn, _, conn_id) in enumerate(connections):
                    if i < recycle_count:
                        try:
                            conn.close()
                            if conn_id in self.connection_last_used:
                                del self.connection_last_used[conn_id]
                            closed_count += 1
                        except Exception as e:
                            logger.error(f"关闭连接时出错: {e}")
                            # 关闭失败，放回池中
                            temp_queue.put(conn)
                    else:
                        # 其余连接放回池中
                        temp_queue.put(conn)
                
                # 将剩余连接放回原池
                while not temp_queue.empty():
                    try:
                        self.pool.put(temp_queue.get(block=False))
                    except queue.Empty:
                        break
                    except queue.Full:
                        # 如果池已满，关闭多余连接
                        try:
                            conn = temp_queue.get(block=False)
                            conn.close()
                            logger.warning("连接池已满，关闭多余连接")
                        except:
                            pass
                
                logger.info(f"已强制回收 {closed_count} 个连接，当前连接池大小: {self.pool.qsize()}")
        
        except Exception as e:
            logger.error(f"强制回收连接时出错: {e}")
    
    def release_connection(self, conn):
        """释放连接回连接池，增加健康检查"""
        try:
            # 检查连接是否健康
            if self._is_connection_healthy(conn):
                # 更新连接最后使用时间
                self.connection_last_used[id(conn)] = time.time()
                
                # 将连接放回池中
                self.pool.put(conn, block=False)
            else:
                # 连接不健康，关闭并创建新连接
                try:
                    conn.close()
                except:
                    pass
                
                # 尝试创建新连接放入池中
                try:
                    new_conn = self._create_connection()
                    self.pool.put(new_conn, block=False)
                except Exception as e:
                    logger.error(f"替换不健康连接失败: {e}")
            
            with self.lock:
                self.active_connections -= 1
        except queue.Full:
            # 如果池已满，关闭连接
            logger.warning("连接池已满，关闭多余连接")
            try:
                conn.close()
            except:
                pass
            
            with self.lock:
                self.active_connections -= 1
        except Exception as e:
            logger.error(f"释放连接时出错: {e}")
            try:
                conn.close()
            except:
                pass
            
            with self.lock:
                self.active_connections -= 1


    def get_stats(self):
        """获取连接池统计信息"""
        with self.lock:
            total_connections = self.pool.qsize() + self.active_connections
            stats = {
                "max_connections": self.max_connections,
                "min_connections": self.min_connections,
                "total_connections": total_connections,
                "active_connections": self.active_connections,
                "available_connections": self.pool.qsize(),
                "usage_percentage": (self.active_connections / max(total_connections, 1)) * 100,
                "avg_usage_rate": sum(self.usage_history) / max(len(self.usage_history), 1)
            }
        return stats
    
    def log_stats(self):
        """记录连接池统计信息到日志"""
        stats = self.get_stats()
        logger.info(f"连接池统计: 总连接={stats['total_connections']}, "
                    f"活跃={stats['active_connections']}, "
                    f"可用={stats['available_connections']}, "
                    f"使用率={stats['usage_percentage']:.1f}%")